import sys
import os
from pathlib import Path
import webbrowser
import base64
import tempfile
import io
import json
import threading
from http.server import HTTPServer, BaseHTTPRequestHandler

# 获取当前文件的绝对路径
current_file = Path(__file__).resolve()
# 计算项目根目录：上溯两级（假设train.py在scripts/目录下）
project_root = current_file.parent.parent
# 将项目根目录添加到系统路径
sys.path.insert(0, str(project_root))
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                            QHBoxLayout, QPushButton, QLabel, QFileDialog, QTabWidget, 
                            QMessageBox, QSpacerItem, QSizePolicy, QGroupBox, QGridLayout, QFrame, QSplitter)
from PyQt5.QtGui import QPixmap, QImage, QFont, QPalette, QColor
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QUrl, pyqtSlot

from PyQt5.QtWebEngineWidgets import QWebEngineView
from PyQt5.QtWebChannel import QWebChannel
from PyQt5.QtCore import QUrl

import torch
import numpy as np
from src.models.vehicle_transformer import VehicleTransformer
from torchvision import transforms

import time
from PIL import Image
import argparse
from tqdm import tqdm
from torch.utils.data import Dataset
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import vit_b_16
from einops import rearrange, repeat
from collections import defaultdict
from torch.utils.data import DataLoader
from torch.optim import AdamW
from pathlib import Path
import matplotlib.pyplot as plt
from collections import Counter
from sklearn.metrics import average_precision_score
from torch.optim.lr_scheduler import CosineAnnealingLR
from src.datasets.veri776_dataset import Veri776Dataset
import warnings
from torch.optim.lr_scheduler import LambdaLR




#线程推理功能
class InferenceThread(QThread):  #用于车辆检索
    finished = pyqtSignal(np.ndarray, list)

    def __init__(self, model, gallery_features, gallery_paths, gallery_camera_ids, 
                 query_img_path, query_camera_id, device):
        super().__init__()
        self.model = model
        self.gallery_features = gallery_features.to(device)
        self.gallery_paths = gallery_paths
        self.gallery_camera_ids = gallery_camera_ids
        self.query_img_path = query_img_path
        self.query_camera_id = query_camera_id
        self.device = device

    def run(self):
        transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        
        img = Image.open(self.query_img_path).convert('RGB')
        img_tensor = transform(img).unsqueeze(0).to(self.device)
        
        with torch.no_grad():
            self.model.eval()
            features = self.model(img_tensor)["bn_feature"]
            features = F.normalize(features, p=2, dim=1)
        
        sim = torch.mm(features, self.gallery_features.T)

        topk = 50
        topk_idx = sim.topk(topk, dim=1)[1].cpu().numpy()[0]

        cross_cam_idx = []
        same_cam_idx = []
        for i in topk_idx:
            if self.gallery_camera_ids[i] != self.query_camera_id:
                cross_cam_idx.append(i)
            else:
                same_cam_idx.append(i)

        combined_idx = cross_cam_idx + same_cam_idx
        result_idx = combined_idx[:6]

        result_paths = [self.gallery_paths[i] for i in result_idx]
        
        self.finished.emit(np.array(img), result_paths)



class InferenceRecognitionThread(QThread):  #用于图片识别
    finished = pyqtSignal(np.ndarray, str, str)  # 发送查询图片、匹配图片路径和识别ID

    def __init__(self, model, gallery_features, gallery_paths, gallery_camera_ids, 
                 query_img_path, query_camera_id, device):
        super().__init__()
        self.model = model
        self.gallery_features = gallery_features.to(device)
        self.gallery_paths = gallery_paths
        self.gallery_camera_ids = gallery_camera_ids
        self.query_img_path = query_img_path
        self.query_camera_id = query_camera_id
        self.device = device

    def run(self):
        transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        
        img = Image.open(self.query_img_path).convert('RGB')
        img_tensor = transform(img).unsqueeze(0).to(self.device)
        
        with torch.no_grad():
            self.model.eval()
            features = self.model(img_tensor)["bn_feature"]
            features = F.normalize(features, p=2, dim=1)
        
        sim = torch.mm(features, self.gallery_features.T)

        topk = 50
        topk_idx = sim.topk(topk, dim=1)[1].cpu().numpy()[0]

        # 找到跨摄像头的匹配项
        cross_cam_idx = []
        for i in topk_idx:
            if self.gallery_camera_ids[i] != self.query_camera_id:
                cross_cam_idx.append(i)
        
        if not cross_cam_idx:
            # 如果没有跨摄像头的匹配项，使用所有匹配项
            result_idx = topk_idx[0]
        else:
            # 使用第一个跨摄像头的匹配项
            result_idx = cross_cam_idx[0]
        
        # 解析识别出的车辆ID
        result_path = self.gallery_paths[result_idx]
        vehicle_id = self.parse_vehicle_id(os.path.basename(result_path))
        
        self.finished.emit(np.array(img), result_path, vehicle_id)
    
    def parse_vehicle_id(self, filename):
        """从文件名中解析车辆ID（格式：0002_c003_...）"""
        try:
            return filename.split('_')[0]
        except:
            return "未知ID"


class InferenceRetrievalThread(QThread):   #用于图片检索
    finished = pyqtSignal(np.ndarray, list, list)  # 发送查询图片、结果路径列表、相似度得分列表

    def __init__(self, model, gallery_features, gallery_paths,  gallery_camera_ids, query_img_path, device):
        super().__init__()
        self.model = model
        self.gallery_features = gallery_features.to(device)
        self.gallery_paths = gallery_paths
        self.gallery_camera_ids = gallery_camera_ids  # 添加摄像头ID列表
        self.query_img_path = query_img_path
        self.device = device

    def run(self):
        transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        
        img = Image.open(self.query_img_path).convert('RGB')
        img_tensor = transform(img).unsqueeze(0).to(self.device)
        
        with torch.no_grad():
            self.model.eval()
            features = self.model(img_tensor)["bn_feature"]
            features = F.normalize(features, p=2, dim=1)
        
        # 计算相似度
        sim = torch.mm(features, self.gallery_features.T)
        sim_scores = sim.cpu().numpy()[0]
        
        # 获取相似度从高到低的索引
        sorted_idx = np.argsort(sim_scores)[::-1]
        
        # 筛选不同摄像头的图片（最多10个）
        result_idx = []
        used_cameras = set()
        max_results = 10
        
        for idx in sorted_idx:
            # 获取该图片的摄像头ID
            camera_id = self.parse_camera_id(self.gallery_paths[idx])
            if camera_id not in used_cameras:
                result_idx.append(idx)
                used_cameras.add(camera_id)
                if len(result_idx) >= max_results:
                    break
        
        result_paths = [self.gallery_paths[i] for i in result_idx]
        result_scores = [f"{sim_scores[i]:.4f}" for i in result_idx]
        
        self.finished.emit(np.array(img), result_paths, result_scores)
    
    def parse_camera_id(self, path):
        """从图片路径中解析摄像头ID"""
        filename = os.path.basename(path)
        try:
            # 格式: 0002_c004_00084250_0 -> c004 -> 4
            cam_part = filename.split('_')[1]
            return int(cam_part[1:])
        except:
            return -1



class VehicleRetrievalWidget(QWidget):        #车辆检索功能，上传图片并获取相似车辆结果
    def __init__(self, main_window, model, gallery_features, gallery_paths, gallery_camera_ids, device):
        super().__init__()
        self.main_window = main_window
        self.model = model
        self.gallery_features = gallery_features
        self.gallery_paths = gallery_paths
        self.gallery_camera_ids = gallery_camera_ids
        self.device = device
        self.current_img_path = None
        self.initUI()
    
    def initUI(self):
        layout = QVBoxLayout()
        layout.setContentsMargins(20, 10, 20, 20)  # 调整整体边距（左、上、右、下）
        layout.setSpacing(15)  # 调整控件间距

        # 标题区域（上移）
        title_container = QWidget()
        title_layout = QVBoxLayout(title_container)
        title_layout.setContentsMargins(0, 5, 0, 5)  # 减小标题区域上下边距
        
        title_label = QLabel("车辆重识别系统")
        title_label.setFont(QFont("SimHei", 16, QFont.Bold))
        title_label.setAlignment(Qt.AlignCenter)
        title_layout.addWidget(title_label)
        
        instruction_label = QLabel("上传车辆图片，系统将在数据库中检索相似车辆")
        instruction_label.setFont(QFont("SimHei", 12))
        instruction_label.setAlignment(Qt.AlignCenter)
        title_layout.addWidget(instruction_label)
        
        layout.addWidget(title_container)

        # 分割线
        line = QFrame()
        line.setFrameShape(QFrame.HLine)
        line.setFrameShadow(QFrame.Sunken)
        layout.addWidget(line)

        # 上传区域
        upload_group = QGroupBox("图片上传")
        upload_layout = QVBoxLayout()
        upload_layout.setContentsMargins(15, 10, 15, 15)  # 调整组框内边距
        
        # 文件名显示
        self.filename_label = QLabel("未选择文件")
        self.filename_label.setAlignment(Qt.AlignCenter)
        self.filename_label.setFont(QFont("SimHei", 11))
        upload_layout.addWidget(self.filename_label)
        
        # 图片预览（增大高度）
        self.preview_label = QLabel("预览区域")
        self.preview_label.setAlignment(Qt.AlignCenter)
        self.preview_label.setMinimumHeight(250)  # 增加预览区域高度
        self.preview_label.setStyleSheet("border: 2px dashed #ddd;")
        upload_layout.addWidget(self.preview_label)
        
        # 按钮布局
        btn_layout = QHBoxLayout()
        btn_layout.setSpacing(20)
        upload_btn = QPushButton("选择文件")
        upload_btn.setFont(QFont("SimHei", 13))
        upload_btn.setMinimumSize(120, 40)  # 增大按钮尺寸
        upload_btn.clicked.connect(self._open_file_dialog)  # 改为调用内部文件对话框方法
        btn_layout.addWidget(upload_btn)
        
        self.search_btn = QPushButton("开始检索")
        self.search_btn.setFont(QFont("SimHei", 13))
        self.search_btn.setMinimumSize(120, 40)
        self.search_btn.clicked.connect(self.start_search)
        self.search_btn.setEnabled(False)
        btn_layout.addWidget(self.search_btn)
        
        upload_layout.addLayout(btn_layout)
        upload_group.setLayout(upload_layout)
        layout.addWidget(upload_group)

        # 结果区域（增大图片尺寸）
        result_group = QGroupBox("检索结果")
        result_group.setFont(QFont("SimHei", 12))
        self.result_layout = QGridLayout()
        self.result_layout.setSpacing(15)  # 增加结果图片间距
        result_group.setLayout(self.result_layout)
        layout.addWidget(result_group)

        # 返回按钮
        back_btn = QPushButton("返回主界面")
        back_btn.setFont(QFont("SimHei", 12))
        back_btn.setMinimumSize(100, 35)
        back_btn.clicked.connect(lambda: self.main_window.tabs.setCurrentIndex(0))
        layout.addWidget(back_btn, alignment=Qt.AlignRight)

        self.setLayout(layout)
    
    def _open_file_dialog(self):
        """内部方法：打开文件对话框选择图片"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择图片", "", "图片文件 (*.jpg *.jpeg *.png *.bmp)"
        )
        if file_path:
            self.upload_image(file_path)  # 调用公共upload_image方法
    

    def upload_image(self, file_path=None):
        """处理图片上传逻辑（支持外部传入路径或文件对话框）"""
        # 如果未传入路径，使用文件对话框获取
        if file_path is None:
            file_path, _ = QFileDialog.getOpenFileName(
                self, "选择图片", "", "图片文件 (*.jpg *.jpeg *.png *.bmp)"
            )
            if not file_path:
                return False
        
        # 检查文件是否存在
        if not os.path.exists(file_path):
            QMessageBox.warning(self, "警告", "文件不存在")
            return False
        
        # 新增：检查文件类型是否为图像
        valid_extensions = [".jpg", ".jpeg", ".png", ".bmp", ".JPG", ".JPEG", ".PNG", ".BMP"]
        file_ext = os.path.splitext(file_path)[1].lower()
        if file_ext not in valid_extensions:
            QMessageBox.warning(self, "警告", "请上传图像文件（.jpg/.jpeg/.png/.bmp）")
            return False  # 非图像文件返回False
        
        # 保存当前图片路径
        self.current_img_path = file_path
        self.filename_label.setText(os.path.basename(file_path))
        

        # 加载并显示预览图
        pixmap = QPixmap(file_path)
        if pixmap.isNull():  # 新增：处理无效图像文件
            QMessageBox.warning(self, "警告", "无法加载该图片，请检查文件格式")
            return False
        
        pixmap = pixmap.scaled(
            400, 250, 
            Qt.KeepAspectRatio, 
            Qt.SmoothTransformation
        )
        self.preview_label.setPixmap(pixmap)
        
        # 启用检索按钮
        self.search_btn.setEnabled(True)
        return True
    
    def start_search(self):
        """开始检索"""
        if not self.current_img_path:
            QMessageBox.warning(self, "警告", "请先选择图片")
            return
        
        filename = os.path.basename(self.current_img_path)
        query_camera_id = self.parse_camera_id(filename)
        
        self.thread = InferenceThread(
            self.model, self.gallery_features, self.gallery_paths,
            self.gallery_camera_ids, self.current_img_path, query_camera_id, self.device
        )
        self.thread.finished.connect(self.display_results)
        self.thread.start()
    
    def display_results(self, img_array, result_paths):
        # 清空之前的结果
        for i in reversed(range(self.result_layout.count())):
            self.result_layout.itemAt(i).widget().deleteLater()
        
        # 转换查询图片（增大显示尺寸）
        img = Image.fromarray(img_array)
        byte_array = io.BytesIO()
        img.save(byte_array, format="JPEG")
        byte_array = byte_array.getvalue()
        
        query_img_label = QLabel("查询图片")
        query_img_label.setFont(QFont("SimHei", 11))
        query_img_label.setAlignment(Qt.AlignCenter)
        query_img = QImage()
        query_img.loadFromData(byte_array)
        pixmap = QPixmap.fromImage(query_img)
        pixmap = pixmap.scaled(
            240, 160,  # 增大查询图片显示尺寸（宽240，高160）
            Qt.KeepAspectRatio, 
            Qt.SmoothTransformation
        )
        query_img_label.setPixmap(pixmap)
        self.result_layout.addWidget(query_img_label, 0, 0)
        
        # 显示检索结果（增大图片尺寸）
        for i, path in enumerate(result_paths):
            if os.path.exists(path):
                result_label = QLabel(f"匹配结果 {i+1}")
                result_label.setFont(QFont("SimHei", 11))
                result_label.setAlignment(Qt.AlignCenter)
                result_pixmap = QPixmap(path)
                result_pixmap = result_pixmap.scaled(
                    240, 160,  # 增大结果图片显示尺寸（宽240，高160）
                    Qt.KeepAspectRatio, 
                    Qt.SmoothTransformation
                )
                result_label.setPixmap(result_pixmap)
                self.result_layout.addWidget(result_label, 0, i+1, 1, 1)  # 调整布局排列
    
    def parse_camera_id(self, filename):
        try:
            cam_part = filename.split('_')[1]
            return int(cam_part[1:])
        except:
            return -1


class ImageRecognitionWidget(QWidget):      #图片识别功能 ，识别图片中的车辆 ID
    def __init__(self, main_window, model, gallery_features, gallery_paths, gallery_camera_ids, device):
        super().__init__()
        self.main_window = main_window
        self.model = model
        self.gallery_features = gallery_features
        self.gallery_paths = gallery_paths
        self.gallery_camera_ids = gallery_camera_ids
        self.device = device
        self.current_img_path = None
        self.initUI()
    
    def initUI(self):
        layout = QVBoxLayout()
        layout.setContentsMargins(20, 20, 20, 20)
        
        # 顶部标题
        title_label = QLabel("请上传车辆图像进行重识别")
        title_label.setFont(QFont("SimHei", 16, QFont.Bold))
        title_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(title_label)
        
        # 分割线
        line = QFrame()
        line.setFrameShape(QFrame.HLine)
        line.setFrameShadow(QFrame.Sunken)
        layout.addWidget(line)
        
        # 主内容区域 - 左右分割
        splitter = QSplitter(Qt.Horizontal)
        
        # 左侧区域 - 输入图片
        left_widget = QWidget()
        left_layout = QVBoxLayout()
        
        input_label = QLabel("输入图片")
        input_label.setFont(QFont("SimHei", 12, QFont.Bold))
        left_layout.addWidget(input_label)
        
        # 输入图片显示区域
        self.input_image_label = QLabel("等待上传图片")
        self.input_image_label.setAlignment(Qt.AlignCenter)
        self.input_image_label.setMinimumHeight(300)
        self.input_image_label.setStyleSheet("border: 2px dashed #ddd;")
        left_layout.addWidget(self.input_image_label)
        
        # 上传图片按钮（蓝底白字）
        upload_btn = QPushButton("上传图片")
        upload_btn.setFont(QFont("SimHei", 14))
        upload_btn.setMinimumSize(150, 40)


        # 使用CSS样式表设置深色背景和白色文字
        upload_btn.setStyleSheet("""
            QPushButton {
                background-color: #0055AA;  /* 更深的蓝色 */
                color: white;
                border-radius: 5px;
            }
            QPushButton:hover {
                background-color: #0077CC;  /* 鼠标悬停时稍微变亮 */
            }
            QPushButton:pressed {
                background-color: #003377;  /* 点击时更深的蓝色 */
            }
        """)


        upload_btn.clicked.connect(self.upload_image)
        left_layout.addWidget(upload_btn, alignment=Qt.AlignCenter)
        
        left_widget.setLayout(left_layout)
        splitter.addWidget(left_widget)
        
        # 右侧区域 - 识别结果
        right_widget = QWidget()
        right_layout = QVBoxLayout()
        
        result_label = QLabel("识别结果")
        result_label.setFont(QFont("SimHei", 12, QFont.Bold))
        right_layout.addWidget(result_label)
        
        # 识别结果图片显示区域
        self.result_image_label = QLabel("等待识别结果")
        self.result_image_label.setAlignment(Qt.AlignCenter)
        self.result_image_label.setMinimumHeight(300)
        self.result_image_label.setStyleSheet("border: 2px solid #ddd;")
        right_layout.addWidget(self.result_image_label)
        
        # 开始识别按钮（蓝底白字）
        self.recognize_btn = QPushButton("开始识别")
        self.recognize_btn.setFont(QFont("SimHei", 14))
        self.recognize_btn.setMinimumSize(150, 40)


       # 使用CSS样式表设置深色背景和白色文字
        self.recognize_btn.setStyleSheet("""
            QPushButton {
                background-color: #0055AA;  /* 更深的蓝色 */
                color: white;
                border-radius: 5px;
            }
            QPushButton:hover {
                background-color: #0077CC;  /* 鼠标悬停时稍微变亮 */
            }
            QPushButton:pressed {
                background-color: #003377;  /* 点击时更深的蓝色 */
            }
        """)




        self.recognize_btn.clicked.connect(self.start_recognition)
        self.recognize_btn.setEnabled(False)
        right_layout.addWidget(self.recognize_btn, alignment=Qt.AlignCenter)
        
        # 识别结果ID显示
        self.id_result_label = QLabel("")
        self.id_result_label.setFont(QFont("SimHei", 14))
        self.id_result_label.setAlignment(Qt.AlignCenter)
        self.id_result_label.setMinimumHeight(40)
        right_layout.addWidget(self.id_result_label)
        
        right_widget.setLayout(right_layout)
        splitter.addWidget(right_widget)
        
        splitter.setSizes([400, 400])  # 初始分割比例
        layout.addWidget(splitter)
        
        # 返回按钮
        back_btn = QPushButton("返回主界面")
        back_btn.setFont(QFont("SimHei", 12))
        back_btn.setMinimumSize(120, 35)
        back_btn.clicked.connect(lambda: self.main_window.tabs.setCurrentIndex(0))
        layout.addWidget(back_btn, alignment=Qt.AlignRight)
        
        self.setLayout(layout)
    
    def upload_image(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择图片", "", "图片文件 (*.jpg *.jpeg *.png *.bmp)"
        )
        if file_path:
            self.current_img_path = file_path
            # 显示上传的图片
            pixmap = QPixmap(file_path)
            pixmap = pixmap.scaled(
                380, 280,  # 调整图片大小以适应显示区域
                Qt.KeepAspectRatio, 
                Qt.SmoothTransformation
            )
            self.input_image_label.setPixmap(pixmap)
            self.input_image_label.setAlignment(Qt.AlignCenter)
            
            # 启用开始识别按钮
            self.recognize_btn.setEnabled(True)
    
    def start_recognition(self):
        if not self.current_img_path:
            QMessageBox.warning(self, "警告", "请先上传图片")
            return
        
        filename = os.path.basename(self.current_img_path)
        query_camera_id = self.parse_camera_id(filename)
        
        self.thread = InferenceRecognitionThread(
            self.model, self.gallery_features, self.gallery_paths,
            self.gallery_camera_ids, self.current_img_path, query_camera_id, self.device
        )
        self.thread.finished.connect(self.display_recognition_results)
        self.thread.start()
    
    def display_recognition_results(self, query_img_array, result_path, vehicle_id):
        # 显示识别结果图片
        if os.path.exists(result_path):
            result_pixmap = QPixmap(result_path)
            result_pixmap = result_pixmap.scaled(
                380, 280,  # 调整图片大小以适应显示区域
                Qt.KeepAspectRatio, 
                Qt.SmoothTransformation
            )
            self.result_image_label.setPixmap(result_pixmap)
            self.result_image_label.setAlignment(Qt.AlignCenter)
        else:
            self.result_image_label.setText("未找到匹配图片")
        
        # 显示识别出的车辆ID
        self.id_result_label.setText(f"系统识别该输入图片的ID为：{vehicle_id}")
    
    def parse_camera_id(self, filename):
        """从文件名中解析摄像头ID（格式：0002_c003_...）"""
        try:
            cam_part = filename.split('_')[1]
            return int(cam_part[1:])
        except:
            return -1



class ImageRetrievalWidget(QWidget):        #图片检索功能 ,查找不同摄像头下的同一车辆
    def __init__(self, main_window, model, gallery_features, gallery_paths,  gallery_camera_ids, device):
        super().__init__()
        self.main_window = main_window
        self.model = model
        self.gallery_features = gallery_features
        self.gallery_paths = gallery_paths
        self.gallery_camera_ids = gallery_camera_ids  # 添加摄像头ID列表
        self.device = device
        self.query_img_path = None
        self.result_paths = []
        self.result_scores = []
        # 显式导入Qt，确保嵌套类可访问
        from PyQt5.QtCore import Qt
        self.Qt = Qt  # 保存到实例属性


        self.initUI()
    
    def initUI(self):
        layout = QVBoxLayout()
        layout.setContentsMargins(20, 20, 20, 20)
        
        # 顶部标题
        title_label = QLabel("图片检索系统")
        title_label.setFont(QFont("SimHei", 18, QFont.Bold))
        title_label.setAlignment(self.Qt.AlignCenter)  # 使用实例属性
        layout.addWidget(title_label)
        
        # 分割线
        line = QFrame()
        line.setFrameShape(QFrame.HLine)
        line.setFrameShadow(QFrame.Sunken)
        layout.addWidget(line)
        
        # 上半部分 - 图片区域
        upper_layout = QHBoxLayout()
        
        # 左侧 - 上传图片区域
        left_widget = QWidget()
        left_layout = QVBoxLayout()
        
        # 待检索图片显示区域
        self.query_image_label = QLabel("等待上传图片")
        self.query_image_label.setAlignment(self.Qt.AlignCenter)  # 使用实例属性
        self.query_image_label.setMinimumSize(300, 200)
        self.query_image_label.setStyleSheet("border: 2px dashed #ddd;")
        left_layout.addWidget(self.query_image_label)
        
        # 上传图片按钮（深蓝色底白字）
        upload_btn = QPushButton("上传图片")
        upload_btn.setFont(QFont("SimHei", 14))
        upload_btn.setMinimumSize(150, 40)
        upload_btn.setStyleSheet("""
            QPushButton {
                background-color: #0055AA;
                color: white;
                border-radius: 5px;
            }
            QPushButton:hover { background-color: #0077CC; }
            QPushButton:pressed { background-color: #003377; }
        """)
        upload_btn.clicked.connect(self.upload_image)
        left_layout.addWidget(upload_btn, alignment=self.Qt.AlignCenter)  # 使用实例属性
        
        left_widget.setLayout(left_layout)
        upper_layout.addWidget(left_widget)
        
        # 右侧 - 检索结果区域
        right_widget = QWidget()
        right_layout = QVBoxLayout()
        
        # 检索结果图片显示区域
        self.result_image_label = QLabel("等待检索结果")
        self.result_image_label.setAlignment(self.Qt.AlignCenter)  # 使用实例属性
        self.result_image_label.setMinimumSize(300, 200)
        self.result_image_label.setStyleSheet("border: 2px solid #ddd;")
        right_layout.addWidget(self.result_image_label)
        
        # 开始检索按钮（深蓝色底白字）
        self.search_btn = QPushButton("开始检索")
        self.search_btn.setFont(QFont("SimHei", 14))
        self.search_btn.setMinimumSize(150, 40)
        self.search_btn.setStyleSheet("""
            QPushButton {
                background-color: #0055AA;
                color: white;
                border-radius: 5px;
            }
            QPushButton:hover { background-color: #0077CC; }
            QPushButton:pressed { background-color: #003377; }
        """)
        self.search_btn.clicked.connect(self.start_search)
        self.search_btn.setEnabled(False)
        right_layout.addWidget(self.search_btn, alignment=self.Qt.AlignCenter)  # 使用实例属性
        
        right_widget.setLayout(right_layout)
        upper_layout.addWidget(right_widget)
        
        layout.addLayout(upper_layout)
        
        # 中间说明文字
        self.tip_label = QLabel("本次的预测结果为上图所示，更多检索结果请点击下方列表查看")
        self.tip_label.setFont(QFont("SimHei", 10))
        self.tip_label.setAlignment(self.Qt.AlignCenter)
        layout.addWidget(self.tip_label)
        
        # 下半部分 - 检索结果列表
        result_group = QGroupBox("检索结果列表")
        result_layout = QVBoxLayout()
        
        # 使用QTableView显示结果列表
        from PyQt5.QtWidgets import QTableView, QHeaderView
        from PyQt5.QtCore import  QAbstractTableModel, QModelIndex
        
        class ResultModel(QAbstractTableModel):
            def __init__(self, data=None, Qt=None):
                super().__init__()
                self.data = data or []
                self.Qt = Qt  # 接收Qt引用
            
            def rowCount(self, parent=QModelIndex()):
                return len(self.data)
            
            def columnCount(self, parent=QModelIndex()):
                return 2
            
            def data(self, index, role):
                if not index.isValid():
                    return None
                if role == self.Qt.DisplayRole:  # 使用传入的Qt
                    return self.data[index.row()][index.column()]
                if role == self.Qt.TextAlignmentRole:  # 使用传入的Qt
                    return Qt.AlignCenter
                return None
            
            def headerData(self, section, orientation, role):
                if role == self.Qt.DisplayRole:  # 使用传入的Qt
                    if orientation == self.Qt.Horizontal:  # 使用传入的Qt
                        return ["图片名称", "相似度得分"][section]
                return None
        
        self.result_model = ResultModel(Qt=self.Qt)
        self.result_table = QTableView()
        self.result_table.setModel(self.result_model)
        self.result_table.setSelectionBehavior(QTableView.SelectRows)
        self.result_table.setEditTriggers(QTableView.NoEditTriggers)
        self.result_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.result_table.clicked.connect(self.on_table_clicked)
        
        result_layout.addWidget(self.result_table)
        result_group.setLayout(result_layout)
        layout.addWidget(result_group)
        
        # 返回按钮
        back_btn = QPushButton("返回主界面")
        back_btn.setFont(QFont("SimHei", 12))
        back_btn.setMinimumSize(120, 35)
        back_btn.clicked.connect(lambda: self.main_window.tabs.setCurrentIndex(0))
        layout.addWidget(back_btn, alignment=self.Qt.AlignRight)  # 使用实例属性
        
        self.setLayout(layout)
    


    def upload_image(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择图片", "", "图片文件 (*.jpg *.jpeg *.png *.bmp)"
        )
        if file_path:
            self.query_img_path = file_path
            # 显示上传的图片
            pixmap = QPixmap(file_path)
            pixmap = pixmap.scaled(
                280, 180, 
                Qt.KeepAspectRatio, 
                Qt.SmoothTransformation
            )
            self.query_image_label.setPixmap(pixmap)
            self.query_image_label.setAlignment(Qt.AlignCenter)
            
            # 启用开始检索按钮
            self.search_btn.setEnabled(True)
    


    def start_search(self):
        if not self.query_img_path:
            QMessageBox.warning(self, "警告", "请先上传图片")
            return
        
        self.thread = InferenceRetrievalThread(
            self.model, self.gallery_features, self.gallery_paths, self.gallery_camera_ids,self.query_img_path, self.device
        )
        self.thread.finished.connect(self.display_retrieval_results)
        self.thread.start()
    



    def display_retrieval_results(self, query_img_array, result_paths, result_scores):
        # 显示查询图片（已在上传时显示）
        
        # 显示第一个检索结果图片
        if result_paths and os.path.exists(result_paths[0]):
            result_pixmap = QPixmap(result_paths[0])
            result_pixmap = result_pixmap.scaled(
                280, 180, 
                self.Qt.KeepAspectRatio, 
                self.Qt.SmoothTransformation
            )
            self.result_image_label.setPixmap(result_pixmap)
            self.result_image_label.setAlignment(self.Qt.AlignCenter)
        else:
            self.result_image_label.setText("未找到检索结果")
        
        # 更新结果列表
        self.result_paths = result_paths
        self.result_scores = result_scores
        self.result_model.data = [[os.path.basename(p), s] for p, s in zip(result_paths, result_scores)]
        self.result_model.layoutChanged.emit()

        # 更新说明文字
        tip_text = f"本次的预测结果为上图所示，共找到{len(result_paths)}个不同摄像头的匹配图片"
        self.tip_label.setText(tip_text)
    


    def on_table_clicked(self, index):
        row = index.row()
        if 0 <= row < len(self.result_paths):
            path = self.result_paths[row]
            if os.path.exists(path):
                pixmap = QPixmap(path)
                pixmap = pixmap.scaled(
                    280, 180, 
                    self.Qt.KeepAspectRatio,  # 使用self.Qt
                    self.Qt.SmoothTransformation  # 使用self.Qt
                )
                self.result_image_label.setPixmap(pixmap)
                self.result_image_label.setAlignment(self.Qt.AlignCenter)




class VehicleReIDGUI(QMainWindow):     #车辆检索主界面
    def __init__(self):
        super().__init__()
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # 加载模型和数据（保持不变）
        self.model = VehicleTransformer(num_classes=776,
            img_size=(224, 224),
            patch_sizes=[16, 8], 
            local_parts=7,
            embed_dim=128,
            depth=4,
            num_heads=4,
            mlp_ratio=4).to(device)
        
        try:
            self.model.load_state_dict(torch.load("checkpoints/best_model.pth", 
                                       map_location=device,
                                       weights_only=False)["state_dict"])
            print("成功加载模型")
        except Exception as e:
            print(f"模型加载失败: {e}")
            self.model.eval()
        
        try:
            gallery_data = torch.load("gallery_features.pt", map_location=device,weights_only=True)
            self.gallery_features = gallery_data["features"].to(device)
            self.gallery_paths = gallery_data["paths"]
            self.gallery_camera_ids = gallery_data.get("camera_ids", [])
            print("成功加载图库数据")
        except Exception as e:
            print(f"图库数据加载失败: {e}")
            self.gallery_features = torch.zeros(0, 128).to(device)
            self.gallery_paths = []
            self.gallery_camera_ids = []
        
        self.device = device
        self.initUI()
        
    def initUI(self):
        self.setWindowTitle("车辆重识别系统")
        self.setGeometry(300, 300, 1600, 900)  # 调整窗口初始尺寸

        self.tabs = QTabWidget()
        self.setCentralWidget(self.tabs)

        # 主检索界面
        main_tab = QWidget()
        main_layout = QVBoxLayout()
        main_layout.setContentsMargins(20, 30, 20, 30)  # 调整主界面边距
        
        main_label = QLabel("主检索界面")
        main_label.setFont(QFont("SimHei", 20, QFont.Bold))  # 增大标题字体
        main_label.setAlignment(Qt.AlignCenter)
        main_layout.addWidget(main_label)
        
        main_layout.addSpacerItem(QSpacerItem(20, 60, QSizePolicy.Minimum, QSizePolicy.Expanding))  # 增加按钮上方空间

        # 主界面按钮布局（增大按钮尺寸）
        button_layout = QHBoxLayout()
        button_layout.setSpacing(40)  # 增加按钮间距
        
        # 车辆检索按钮
        web_btn = QPushButton("车辆检索")
        web_btn.setFont(QFont("SimHei", 16))
        web_btn.setMinimumSize(220, 70)  # 增大按钮尺寸（宽220，高70）
        web_btn.clicked.connect(lambda: self.tabs.setCurrentIndex(1))
        button_layout.addWidget(web_btn)
        
       
        
        # 图片识别按钮
        recog_btn = QPushButton("图片识别")
        recog_btn.setFont(QFont("SimHei", 16))
        recog_btn.setMinimumSize(220, 70)
        recog_btn.clicked.connect(lambda: self.tabs.setCurrentIndex(2))
        button_layout.addWidget(recog_btn)
        
        # 图片检索按钮
        retri_btn = QPushButton("图片检索")
        retri_btn.setFont(QFont("SimHei", 16))
        retri_btn.setMinimumSize(220, 70)
        retri_btn.clicked.connect(lambda: self.tabs.setCurrentIndex(3))
        button_layout.addWidget(retri_btn)
        
        main_layout.addLayout(button_layout)
        main_layout.addSpacerItem(QSpacerItem(20, 60, QSizePolicy.Minimum, QSizePolicy.Expanding))  # 增加按钮下方空间
        
        main_tab.setLayout(main_layout)
        self.tabs.addTab(main_tab, "主检索界面")

        # 添加功能选项卡（保持不变）
        self.tabs.addTab(
            VehicleRetrievalWidget(self, self.model, self.gallery_features, self.gallery_paths, self.gallery_camera_ids, self.device),
            "车辆检索"
        )
       


        self.tabs.addTab(
            ImageRecognitionWidget(self, self.model, self.gallery_features, self.gallery_paths, self.gallery_camera_ids, self.device),
            "图片识别"
        )

        # 添加功能选项卡（修改图片检索选项卡）
        self.tabs.addTab(
            ImageRetrievalWidget(self, self.model, self.gallery_features, self.gallery_paths,self.gallery_camera_ids,  self.device),
            "图片检索"
        )

        font = QFont("SimHei", 12)  # 增大选项卡字体
        self.tabs.setFont(font)



    def upload_image(self, file_path):
        """公开上传图片的方法（用于测试）"""
        # 获取车辆检索选项卡中的Widget
        if self.tabs.currentIndex() != 1:
            self.tabs.setCurrentIndex(1)  # 切换到车辆检索选项卡
        
        retrieval_widget = self.tabs.widget(1)
        if hasattr(retrieval_widget, 'upload_image'):
            # 调用VehicleRetrievalWidget的upload_image方法
            return retrieval_widget.upload_image(file_path)
        return False


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = VehicleReIDGUI()
    window.show()
    sys.exit(app.exec_())


